package org.example.search.service.impl;

import com.alibaba.fastjson2.JSON;
import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.DropCollectionReq;
import io.milvus.v2.service.collection.request.GetLoadStateReq;
import io.milvus.v2.service.index.request.CreateIndexReq;
import io.milvus.v2.service.index.request.DropIndexReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.UpsertReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import io.milvus.v2.service.vector.response.UpsertResp;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.example.search.entity.SearchDto;
import org.example.search.service.MilvusService;
import org.example.search.util.BGEM3ApiClient;
import org.example.search.util.GsonUtils;
import org.springframework.stereotype.Service;

import java.util.*;

@Slf4j
@Service
public class MilvusServiceImpl implements MilvusService {

    @Resource
    private MilvusClientV2 client;

    private static final String COLLECTION_NAME = "product_setup_1";

    @Override
    public Object init() {
        // 创建schema
        CreateCollectionReq.CollectionSchema schema = client.createSchema();

        schema.addField(AddFieldReq.builder()
                .fieldName("productId")
                .dataType(DataType.Int64)
                .isPrimaryKey(true)
                .autoID(false)
                .build());

        schema.addField(AddFieldReq.builder()
                .fieldName("productVector")
                .dataType(DataType.FloatVector)
                .dimension(1024)
                .build());

        schema.addField(AddFieldReq.builder()
                .fieldName("productTitle")
                .dataType(DataType.VarChar)
                .maxLength(200)
                .build());
        //  准备索引参数

        IndexParam indexParamForVectorField = IndexParam.builder()
                .fieldName("productVector")
                .indexType(IndexParam.IndexType.IVF_FLAT)
                .metricType(IndexParam.MetricType.COSINE)
                .build();

        List<IndexParam> indexParams = new ArrayList<>();
        indexParams.add(indexParamForVectorField);

        //  Create a collection with schema and index parameters
        CreateCollectionReq customizedSetupReq1 = CreateCollectionReq.builder()
                .collectionName(COLLECTION_NAME)
                .collectionSchema(schema)
                .indexParams(indexParams)
                .build();

        client.dropCollection(DropCollectionReq.builder().collectionName(COLLECTION_NAME).build());
        client.createCollection(customizedSetupReq1);

        // 3.5 Get load state of the collection
        GetLoadStateReq customSetupLoadStateReq1 = GetLoadStateReq.builder()
                .collectionName(COLLECTION_NAME)
                .build();

        Boolean loaded = client.getLoadState(customSetupLoadStateReq1);
        log.info("创建集合：{}", loaded);
        return "success";
    }

    @Override
    public InsertResp insert(SearchDto searchDto) throws Exception {
        JsonObject jsonObject = GsonUtils.objectToJsonObject(searchDto);
        jsonObject.add("productVector", BGEM3ApiClient.getEmbedding(searchDto.getProductTitle()));
        List<JsonObject> data = Collections.singletonList(jsonObject);

        InsertReq insertReq = InsertReq.builder()
                .collectionName(COLLECTION_NAME)
                .data(data)
                .build();

        return client.insert(insertReq);
    }

    @Override
    public UpsertResp upsert(SearchDto searchDto) throws Exception {
        JsonObject jsonObject = GsonUtils.objectToJsonObject(searchDto);
        jsonObject.add("productVector", BGEM3ApiClient.getEmbedding(searchDto.getProductTitle()));

        List<JsonObject> data = new ArrayList<>();
        data.add(jsonObject);

        UpsertReq upsertReq = UpsertReq.builder()
                .collectionName(COLLECTION_NAME)
                .data(data)
                .build();
        return client.upsert(upsertReq);
    }

    @Override
    public SearchResp searchVec(String title) throws Exception {
        FloatVec queryVector = new FloatVec(BGEM3ApiClient.getEmbeddingList(title));
        SearchReq searchReq = SearchReq.builder()
                .collectionName(COLLECTION_NAME)
                .data(Collections.singletonList(queryVector))
                .outputFields(Collections.singletonList("productTitle"))
                .topK(10)
                .build();

        return client.search(searchReq);

    }

    @Override
    public Object createIndex() {
        IndexParam indexParamForVectorField = IndexParam.builder()
                .fieldName("productVector")
                .indexType(IndexParam.IndexType.IVF_FLAT)
                .metricType(IndexParam.MetricType.COSINE)
                .build();

        CreateIndexReq createIndexReq = CreateIndexReq.builder()
                .collectionName(COLLECTION_NAME)
                .indexParams(Collections.singletonList(indexParamForVectorField))
                .build();
        client.createIndex(createIndexReq);
        return "索引创建成功";
    }

    @Override
    public Object dropIndex() {
        DropIndexReq dropIndexReq = DropIndexReq.builder()
                .collectionName(COLLECTION_NAME)
                .fieldName("productVector")
                .build();
        client.dropIndex(dropIndexReq);
        return "索引删除成功";
    }

    // 以下方法仅用于测试
    protected JsonArray getEmbedding(String text) throws Exception {
        return BGEM3ApiClient.getEmbedding(text).getAsJsonArray();
    }

    protected List<Float> getEmbeddingList(String text) throws Exception {
        float[] embeddingArray = BGEM3ApiClient.getEmbeddingList(text);
        List<Float> embeddingList = new ArrayList<>();
        for (float value : embeddingArray) {
            embeddingList.add(value);
        }
        return embeddingList;
    }
} 